from pathlib import Path

from src.module.Embedding import Embedding
from src.constance import SEARCH_ARTICLE_TEXT, db_file_path

from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

class ChromaStorage:
    def __init__(self):
        self.embedding = Embedding().embedding
        self.local_db_file = str(db_file_path('test-1_db'))

    @staticmethod
    def get_text_documents(docs):
        return RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(docs)

    @staticmethod
    def parser_file():
        return TextLoader(SEARCH_ARTICLE_TEXT, encoding="utf8").load()

    def get_docs(self):
        return self.get_text_documents(self.parser_file())

    def search(self):
        docs = self.get_docs()
        vector = Chroma.from_documents(docs, self.embedding)
        print('------------>', vector.similarity_search('第一个故事是什么?'))

    def search_local(self):
        docs = self.get_docs()

        if Path(self.local_db_file).exists():
            vector = Chroma(persist_directory=self.local_db_file, embedding_function=self.embedding)
            print('---> 加载向量数据库...')
        else:
            vector = Chroma.from_documents(docs, self.embedding, persist_directory=self.local_db_file)

        # 直接采用原始方式查询
        # print('------------>', vector.similarity_search('第二个故事是什么?')) ## 查询相似文本
        # print('------------> 2', vector.similarity_search_with_score('第二个故事是什么?')) ## 查询相似文本的余弦相似度

        # 采用 retriever 查询，推荐使用这种方式查询
        retriever = vector.as_retriever(
            search_type="mmr", ## 将分块内容查询出来后整合一次在进行分析，得出答案
            search_kwargs={'k': 5} ## 仅查询前 5 条
        )
        print(retriever.invoke('第二个故事是什么?'))

    def start(self):
        self.search_local()